#ifndef WEIGHTED_RANDOM_HPP
#define WEIGHTED_RANDOM_HPP

#include <random>
#include <algorithm>

#include "./LoadBalancePolicyIface.h"

namespace LoadBalance
{
    template <typename T>
    class WeightedRandom : public LoadBalancePolicyIface<T>
    {
    public:
        explicit WeightedRandom(const std::vector<std::pair<T, int>> &nodes) : nodes_(nodes)
        {
            for (const auto &node : this->nodes_)
            {
                this->totalWeight_ += node.second;
                this->cumulativeWeights_.push_back(this->totalWeight_);
            }
        }

        T getNext() override
        {
            std::lock_guard<std::mutex> lock(this->mtx_);

            if (this->nodes_.empty()) // 注意错误处理
            {
                return T{};
            }

            // 有 节点权重/总权重 的概率选中，[1,总权重]的随机数有更大概率落在大区间并选择该节点
            // 如1 2 3，总权重6，随机数1会选择1节点，23会选择2节点，456会选择3节点
            // 每次累加权重判断概率区间效率低->提前计算累积权重/前缀和+二分查找获取
            std::uniform_int_distribution<int> dist(1, this->totalWeight_); // [1,totalWeight]随机数
            int randomValue = dist(this->randomEngine_);

            // int cumulativeWeight = 0;
            // for (const auto &node : this->nodes_)
            // {
            //     cumulativeWeight += node.second;     // 累积权重
            //     if (randomValue <= cumulativeWeight) // 随机值落在该节点权重区间
            //     {
            //         return node.first;
            //     }
            // }
            // return T{}; // 不会到这

            auto it = std::upper_bound(this->cumulativeWeights_.begin(), this->cumulativeWeights_.end(), randomValue); // 第一个大于randomValue的累积权重
            size_t idx = it - this->cumulativeWeights_.begin() - 1;                                                    // -1就是最后一个小于randomValue的累积权重
            return this->nodes_[idx].first;
        }

    private:
        std::vector<std::pair<T, int>> nodes_; // 节点-权重
        int totalWeight_ = 0;                  // 总权重
        std::vector<int> cumulativeWeights_;   // 每个节点的累积权重/前缀和，用于二分查找概率区间

        std::mt19937 randomEngine_; // 随机数引擎/生成器
        std::mutex mtx_;            // 注意线程安全
    };
} // namespace LoadBalance

#endif // WEIGHTED_RANDOM_HPP